import os

def get_path(use_mydata,is_test,path_root_big='classification_24category',path_root_my='mydata'): 
    """获得路径
    
    Args:
        is_test       :是否处于测试状态
        use_mydata    :是否使用自己的数据集
        path_root_big :异源数据集目录
        path_root_my  :自己的少量数据集目录，暂时不存在
    Returns:
        path_txt      :数据记录的位置
        path_root     :数据集目录（train/test）
    """

    name_mode = ['train','test'][is_test]
    
    # 数据集目录（train/test的上一级）
    path_root_dataset = path_root_my if use_mydata else path_root_big
    
    path_txt = f'./{path_root_dataset}/label_{name_mode}.txt' 
    path_root = f'./{path_root_dataset}/{name_mode}'

    return (path_txt,path_root)

def load_data(path_txt,path_root):
    """生成记录文件
    
    
    path_txt : 记录文件路径
    path_root: 数据集路径
    """
    print(f'{path_root} -> {path_txt}')
    assert(os.path.exists(path_root))
    # 计数
    count =[0]*24
    # 删除文件
    if(os.path.exists(path_txt)):
        os.remove(path_txt)
    with open(path_txt,'w') as f:
        # 文件夹list，24种染色体图片
        path_dir_name_s = list(range(1,23)) +['X','Y'] # 24种染色体的文件夹名
        for i,path_name_dir in enumerate(path_dir_name_s):
            # 某一种染色体的文件夹
            path_dir = f'{path_root}/{path_name_dir}'
            path_list = os.listdir(path_dir)
            count[i] = len(path_list)
            for path_img in path_list:
                # 每一张染色体图片的路径
                path_img = f'{path_dir}/{path_img}'
                # 输出一行（路径）+种类
                f.write(f'{path_img} {i}\n')
    print(sum(count),count)
    return count